import importlib
import os
from typing import Dict, Iterable, List, Optional, Any


AI_PROVIDER = "qwen"
AI_MODEL = "Qwen3-Instruct"

QWEN_API_KEY = None
QWEN_BASE_URL = "http://192.168.36.201:9997/v1"

OLLAMA_API_KEY = "ollama"
OLLAMA_BASE_URL = "http://localhost:11434/v1"
OLLAMA_MODEL = "llama3.1"

LMSTUDIO_API_KEY = "lm-studio"
LMSTUDIO_BASE_URL = "http://localhost:1234/v1"
LMSTUDIO_MODEL = "qwen/qwen3-8b"


PROVIDERS: Dict[str, str] = {
    "qwen": "process.ai.qwen_local",
    "qwenlocal": "process.ai.qwen_local",
    "qwencloud": "process.ai.qwen",
    "ollama": "process.ai.ollama",
    "lmstudio": "process.ai.lmstudio",
}


def _get_module(name: Optional[str] = None):
    provider = (name or AI_PROVIDER or "qwen").lower().strip().replace("-", "")
    module_path = PROVIDERS.get(provider)
    if not module_path:
        raise RuntimeError(f"Unsupported AI provider: {provider}")
    return importlib.import_module(module_path)


def _provider_config(provider: str) -> Dict[str, Optional[str]]:
    p = provider.lower().strip()
    if p in ("qwen", "qwenlocal"):
        return {"api_key": QWEN_API_KEY, "base_url": QWEN_BASE_URL, "model": "Qwen3-Instruct"}
    if p == "qwencloud":
        return {"api_key": os.getenv("DASHSCOPE_API_KEY"), "base_url": os.getenv("DASHSCOPE_BASE_URL"), "model": "qwen-plus"}
    if p == "ollama":
        return {"api_key": OLLAMA_API_KEY, "base_url": OLLAMA_BASE_URL, "model": OLLAMA_MODEL}
    if p == "lmstudio":
        return {"api_key": LMSTUDIO_API_KEY, "base_url": LMSTUDIO_BASE_URL, "model": LMSTUDIO_MODEL}
    return {"api_key": None, "base_url": None, "model": "qwen-plus"}


def get_client(
    provider: Optional[str] = None,
    api_key: Optional[str] = None,
    base_url: Optional[str] = None,
):
    mod = _get_module(provider)
    prov = (provider or AI_PROVIDER or "qwen")
    cfg = _provider_config(prov)
    key = api_key if api_key is not None else cfg["api_key"]
    url = base_url if base_url is not None else cfg["base_url"]
    return mod.get_client(api_key=key, base_url=url)


def chat(
    messages: List[Dict[str, str]],
    model: Optional[str] = None,
    stream: bool = False,
    provider: Optional[str] = None,
    **kwargs,
) -> str | Iterable[str]:
    mod = _get_module(provider)
    prov = (provider or AI_PROVIDER or "qwen")
    cfg = _provider_config(prov)
    get_client(provider=prov, api_key=cfg["api_key"], base_url=cfg["base_url"])  
    m = (model or AI_MODEL or cfg["model"]) or "qwen-plus"
    return mod.chat(messages=messages, model=str(m), stream=stream, **kwargs)


def chat_with_usage(
    messages: List[Dict[str, str]],
    model: Optional[str] = None,
    provider: Optional[str] = None,
    **kwargs,
) -> Dict[str, Any]:
    mod = _get_module(provider)
    prov = (provider or AI_PROVIDER or "qwen")
    cfg = _provider_config(prov)
    get_client(provider=prov, api_key=cfg["api_key"], base_url=cfg["base_url"])
    m = (model or AI_MODEL or cfg["model"]) or "qwen-plus"
    return mod.chat_with_usage(messages=messages, model=str(m), **kwargs)


def simple_chat(
    prompt: str,
    system: Optional[str] = None,
    model: Optional[str] = None,
    stream: bool = False,
    provider: Optional[str] = None,
    **kwargs,
) -> str | Iterable[str]:
    mod = _get_module(provider)
    prov = (provider or AI_PROVIDER or "qwen")
    cfg = _provider_config(prov)
    get_client(provider=prov, api_key=cfg["api_key"], base_url=cfg["base_url"])  
    m = (model or AI_MODEL or cfg["model"]) or "qwen-plus"
    return mod.simple_chat(prompt=prompt, system=system, model=str(m), stream=stream, **kwargs)


def simple_chat_with_usage(
    prompt: str,
    system: Optional[str] = None,
    model: Optional[str] = None,
    provider: Optional[str] = None,
    **kwargs,
) -> Dict[str, Any]:
    mod = _get_module(provider)
    prov = (provider or AI_PROVIDER or "qwen")
    cfg = _provider_config(prov)
    get_client(provider=prov, api_key=cfg["api_key"], base_url=cfg["base_url"]) 
    m = (model or AI_MODEL or cfg["model"]) or "qwen-plus"
    return mod.simple_chat_with_usage(prompt=prompt, system=system, model=str(m), **kwargs)


def _default_model(mod) -> str:
    name = getattr(mod, "__name__", "")
    if name.endswith("qwen"):
        return "qwen-plus"
    if name.endswith("ollama"):
        return OLLAMA_MODEL
    if name.endswith("lmstudio"):
        return LMSTUDIO_MODEL
    return "qwen-plus"


__all__ = [
    "get_client",
    "chat",
    "chat_with_usage",
    "simple_chat",
    "simple_chat_with_usage",
]
